package com.jorado.metrics.jvm;

import com.jorado.metrics.util.LabelHelp;
import com.jorado.metrics.util.VersionHelp;
import io.prometheus.client.Collector;
import io.prometheus.client.GaugeMetricFamily;

import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryPoolMXBean;
import java.lang.management.MemoryUsage;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * Exports metrics about JVM memory areas.
 * <p>
 * Example usage:
 * <pre>
 * {@code
 *   new MemoryPoolsExports().register();
 * }
 * </pre>
 * Example metrics being exported:
 * <pre>
 *   jvm_memory_bytes_used{area="heap"} 2000000
 *   jvm_memory_bytes_committed{area="nonheap"} 200000
 *   jvm_memory_bytes_max{area="nonheap"} 2000000
 *   jvm_memory_pool_bytes_used{pool="PS Eden Space"} 2000
 * </pre>
 */
public class MemoryPoolsExports extends Collector {
    private final MemoryMXBean memoryBean;
    private final List<MemoryPoolMXBean> poolBeans;

    public MemoryPoolsExports() {
        this(
                ManagementFactory.getMemoryMXBean(),
                ManagementFactory.getMemoryPoolMXBeans());
    }

    public MemoryPoolsExports(MemoryMXBean memoryBean,
                              List<MemoryPoolMXBean> poolBeans) {
        this.memoryBean = memoryBean;
        this.poolBeans = poolBeans;
    }

    void addMemoryAreaMetrics(List<MetricFamilySamples> sampleFamilies) {
        MemoryUsage heapUsage = memoryBean.getHeapMemoryUsage();
        MemoryUsage nonHeapUsage = memoryBean.getNonHeapMemoryUsage();
        GaugeMetricFamily used = new GaugeMetricFamily(
                "jvm_memory_bytes_used",
                "Used bytes of a given JVM memory area." + VersionHelp.Ver,
                Arrays.asList(LabelHelp.addLabels("area")));
        used.addMetric(Arrays.asList(LabelHelp.addValues("heap")), heapUsage.getUsed());
        used.addMetric(Arrays.asList(LabelHelp.addValues("nonheap")), nonHeapUsage.getUsed());
        sampleFamilies.add(used);
        GaugeMetricFamily committed = new GaugeMetricFamily(
                "jvm_memory_bytes_committed",
                "Committed (bytes) of a given JVM memory area." + VersionHelp.Ver,
                Arrays.asList(LabelHelp.addLabels("area")));
        committed.addMetric(Arrays.asList(LabelHelp.addValues("heap")), heapUsage.getCommitted());
        committed.addMetric(Arrays.asList(LabelHelp.addValues("nonheap")), nonHeapUsage.getCommitted());
        sampleFamilies.add(committed);
        GaugeMetricFamily max = new GaugeMetricFamily(
                "jvm_memory_bytes_max",
                "Max (bytes) of a given JVM memory area." + VersionHelp.Ver,
                Arrays.asList(LabelHelp.addLabels("area")));
        max.addMetric(Arrays.asList(LabelHelp.addValues("heap")), heapUsage.getMax());
        max.addMetric(Arrays.asList(LabelHelp.addValues("nonheap")), nonHeapUsage.getMax());
        sampleFamilies.add(max);
        GaugeMetricFamily init = new GaugeMetricFamily(
                "jvm_memory_bytes_init",
                "Initial bytes of a given JVM memory area." + VersionHelp.Ver,
                Arrays.asList(LabelHelp.addLabels("area")));
        init.addMetric(Arrays.asList(LabelHelp.addValues("heap")), heapUsage.getInit());
        init.addMetric(Arrays.asList(LabelHelp.addValues("nonheap")), nonHeapUsage.getInit());
        sampleFamilies.add(init);
    }

    void addMemoryPoolMetrics(List<MetricFamilySamples> sampleFamilies) {
        GaugeMetricFamily used = new GaugeMetricFamily(
                "jvm_memory_pool_bytes_used",
                "Used bytes of a given JVM memory pool." + VersionHelp.Ver,
                Arrays.asList(LabelHelp.addLabels("area")));
        sampleFamilies.add(used);
        GaugeMetricFamily committed = new GaugeMetricFamily(
                "jvm_memory_pool_bytes_committed",
                "Committed bytes of a given JVM memory pool." + VersionHelp.Ver,
                Arrays.asList(LabelHelp.addLabels("area")));
        sampleFamilies.add(committed);
        GaugeMetricFamily max = new GaugeMetricFamily(
                "jvm_memory_pool_bytes_max",
                "Max bytes of a given JVM memory pool." + VersionHelp.Ver,
                Arrays.asList(LabelHelp.addLabels("area")));
        sampleFamilies.add(max);
        GaugeMetricFamily init = new GaugeMetricFamily(
                "jvm_memory_pool_bytes_init",
                "Initial bytes of a given JVM memory pool." + VersionHelp.Ver,
                Arrays.asList(LabelHelp.addLabels("area")));
        sampleFamilies.add(init);
        for (final MemoryPoolMXBean pool : poolBeans) {
            MemoryUsage poolUsage = pool.getUsage();
            used.addMetric(
                    Arrays.asList(LabelHelp.addValues(pool.getName())),
                    poolUsage.getUsed());
            committed.addMetric(
                    Arrays.asList(LabelHelp.addValues(pool.getName())),
                    poolUsage.getCommitted());
            max.addMetric(
                    Arrays.asList(LabelHelp.addValues(pool.getName())),
                    poolUsage.getMax());
            init.addMetric(
                    Arrays.asList(LabelHelp.addValues(pool.getName())),
                    poolUsage.getInit());
        }
    }

    public List<MetricFamilySamples> collect() {
        List<MetricFamilySamples> mfs = new ArrayList<MetricFamilySamples>();
        addMemoryAreaMetrics(mfs);
        addMemoryPoolMetrics(mfs);
        return mfs;
    }
}
